import os
import torch
import torch.nn as nn 
import torch.nn.functional as F
import struct


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.relu = nn.ReLU()
    
    def forward(self, x1, x2):
        out = torch.add(x1, x2)
        out = self.relu(out)
        return out
    
def export(model, onnx_path):
    input1 = torch.randn(1, 20)
    input2 = torch.randn(1, 20)
    torch.onnx.export(model, (input1, input2), onnx_path, 
        verbose=True,
        opset_version=11, 
        input_names=['x1', 'x2'], 
        output_names=['y'], 
    )


if __name__ == '__main__':
    model = Model()
    onnx_path = './add.onnx'
    export(model, onnx_path)

